#include "sm3.h"

#include <string.h>

// SM3_ENDIAN
#define SM3_LITTLEENDIAN  0   //!< Little endian machine
#define SM3_BIGENDIAN     1   //!< Big endian machine

// Endianness of the machine.
#ifndef SM3_ENDIAN
// Detect with GCC 4.6's macro
#  ifdef __BYTE_ORDER__
#    if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
#      define SM3_ENDIAN SM3_LITTLEENDIAN
#    elif __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
#      define SM3_ENDIAN SM3_BIGENDIAN
#    else
#      error Unknown machine endianess detected. User needs to define SM3_ENDIAN.
#    endif // __BYTE_ORDER__
// Detect with GLIBC's endian.h
#  elif defined(__GLIBC__)
#    include <endian.h>
#    if (__BYTE_ORDER == __LITTLE_ENDIAN)
#      define SM3_ENDIAN SM3_LITTLEENDIAN
#    elif (__BYTE_ORDER == __BIG_ENDIAN)
#      define SM3_ENDIAN SM3_BIGENDIAN
#    else
#      error Unknown machine endianess detected. User needs to define SM3_ENDIAN.
#   endif // __GLIBC__
// Detect with _LITTLE_ENDIAN and _BIG_ENDIAN macro
#  elif defined(_LITTLE_ENDIAN) && !defined(_BIG_ENDIAN)
#    define SM3_ENDIAN SM3_LITTLEENDIAN
#  elif defined(_BIG_ENDIAN) && !defined(_LITTLE_ENDIAN)
#    define SM3_ENDIAN SM3_BIGENDIAN
// Detect with architecture macros
#  elif defined(__sparc) || defined(__sparc__) || defined(_POWER) || defined(__powerpc__) || defined(__ppc__) || defined(__hpux) || defined(__hppa) || defined(_MIPSEB) || defined(_POWER) || defined(__s390__)
#    define SM3_ENDIAN SM3_BIGENDIAN
#  elif defined(__i386__) || defined(__alpha__) || defined(__ia64) || defined(__ia64__) || defined(_M_IX86) || defined(_M_IA64) || defined(_M_ALPHA) || defined(__amd64) || defined(__amd64__) || defined(_M_AMD64) || defined(__x86_64) || defined(__x86_64__) || defined(_M_X64) || defined(__bfin__)
#    define SM3_ENDIAN SM3_LITTLEENDIAN
#  elif defined(_MSC_VER) && defined(_M_ARM)
#    define SM3_ENDIAN SM3_LITTLEENDIAN
#  elif defined(RAPIDJSON_DOXYGEN_RUNNING)
#    define SM3_ENDIAN
#  else
#    error Unknown machine endianess detected. User needs to define SM3_ENDIAN.   
#  endif
#endif // SM3_ENDIAN

#if SM3_ENDIAN	// Big endian
#  define sm3_cpu_to_be16(v) (v)
#  define sm3_cpu_to_be32(v) (v)
#  define be16_cpu_to_cpu(v) (v)
#  define be32_cpu_to_cpu(v) (v)
#else // Little endian
#  define sm3_cpu_to_be16(v) (((v) << 8) | ((v) >> 8))
#  define sm3_cpu_to_be32(v) (((v) >> 24) | (((v) >> 8) & 0xff00) | (((v) << 8) & 0xff0000) | ((v) << 24))
#  define be16_cpu_to_cpu(v) sm3_cpu_to_be16(v)
#  define be32_cpu_to_cpu(v) sm3_cpu_to_be32(v)
#endif

#define SM3_ROTATELEFT(X, n)  (((X) << (n)) | ((X) >> (32 - (n))))

#define SM3_P0(x) ((x) ^ SM3_ROTATELEFT((x), 9) ^ SM3_ROTATELEFT((x), 17))
#define SM3_P1(x) ((x) ^ SM3_ROTATELEFT((x), 15) ^ SM3_ROTATELEFT((x), 23))

#define SM3_FF0(x, y, z) ((x) ^ (y) ^ (z))
#define SM3_FF1(x, y, z) (((x) & (y)) | ((x) & (z)) | ((y) & (z)))

#define SM3_GG0(x, y, z) ((x) ^ (y) ^ (z))
#define SM3_GG1(x, y, z) (((x) & (y)) | ((~(x)) & (z)))

void sm3_init(sm3_ctx_t* ctx)
{
    ctx->digest[0] = 0x7380166F;
    ctx->digest[1] = 0x4914B2B9;
    ctx->digest[2] = 0x172442D7;
    ctx->digest[3] = 0xDA8A0600;
    ctx->digest[4] = 0xA96F30BC;
    ctx->digest[5] = 0x163138AA;
    ctx->digest[6] = 0xE38DEE4D;
    ctx->digest[7] = 0xB0FB0E4E;

    ctx->nblocks = 0;
    ctx->num = 0;
}

void sm3_update(sm3_ctx_t* ctx, const unsigned char* data, size_t data_len)
{
    if (ctx->num) {
        unsigned int left = SM3_BLOCK_SIZE - ctx->num;
        if (data_len < left) {
            memcpy(ctx->block + ctx->num, data, data_len);
            ctx->num += data_len;
            return;
        } else {
            memcpy(ctx->block + ctx->num, data, left);
            sm3_compress(ctx->digest, ctx->block);
            ctx->nblocks++;
            data += left;
            data_len -= left;
        }
    }
    while (data_len >= SM3_BLOCK_SIZE) {
        sm3_compress(ctx->digest, data);
        ctx->nblocks++;
        data += SM3_BLOCK_SIZE;
        data_len -= SM3_BLOCK_SIZE;
    }
    ctx->num = data_len;
    if (data_len) {
        memcpy(ctx->block, data, data_len);
    }
}

void sm3_final(sm3_ctx_t *ctx, unsigned char *digest)
{
	uint32_t *pdigest = (uint32_t*)digest;
	uint32_t *count = (uint32_t*)(ctx->block + SM3_BLOCK_SIZE - 8);

	ctx->block[ctx->num] = 0x80;

	if (ctx->num + 9 <= SM3_BLOCK_SIZE) {
		memset(ctx->block + ctx->num + 1, 0, SM3_BLOCK_SIZE - ctx->num - 9);
	} else {
		memset(ctx->block + ctx->num + 1, 0, SM3_BLOCK_SIZE - ctx->num - 1);
		sm3_compress(ctx->digest, ctx->block);
		memset(ctx->block, 0, SM3_BLOCK_SIZE - 8);
	}

	count[0] = sm3_cpu_to_be32((ctx->nblocks) >> 23);
	count[1] = sm3_cpu_to_be32((ctx->nblocks << 9) + (ctx->num << 3));

	sm3_compress(ctx->digest, ctx->block);
	for (int i = 0; i < sizeof(ctx->digest) / sizeof(ctx->digest[0]); ++i) {
		pdigest[i] = sm3_cpu_to_be32(ctx->digest[i]);
	}
}

void sm3_compress(uint32_t digest[8], const unsigned char block[64])
{
	int j;
	uint32_t W[68], W1[64];
	const uint32_t* pblock = (const uint32_t*)block;

	uint32_t A = digest[0];
	uint32_t B = digest[1];
	uint32_t C = digest[2];
	uint32_t D = digest[3];
	uint32_t E = digest[4];
	uint32_t F = digest[5];
	uint32_t G = digest[6];
	uint32_t H = digest[7];
	uint32_t SS1, SS2, TT1, TT2, T[64];

	for (j = 0; j < 16; ++j) {
		W[j] = sm3_cpu_to_be32(pblock[j]);
	}
	for (j = 16; j < 68; ++j) {
		W[j] = SM3_P1(W[j - 16] ^ W[j - 9] ^ SM3_ROTATELEFT(W[j - 3], 15)) ^ SM3_ROTATELEFT(W[j - 13], 7) ^ W[j - 6];;
	}
	for (j = 0; j < 64; ++j) {
		W1[j] = W[j] ^ W[j + 4];
	}

	for (j = 0; j < 16; ++j) {
		T[j] = 0x79CC4519;
		SS1 = SM3_ROTATELEFT((SM3_ROTATELEFT(A, 12) + E + SM3_ROTATELEFT(T[j], j)), 7);
		SS2 = SS1 ^ SM3_ROTATELEFT(A, 12);
		TT1 = SM3_FF0(A, B, C) + D + SS2 + W1[j];
		TT2 = SM3_GG0(E, F, G) + H + SS1 + W[j];
		D = C;
		C = SM3_ROTATELEFT(B, 9);
		B = A;
		A = TT1;
		H = G;
		G = SM3_ROTATELEFT(F, 19);
		F = E;
		E = SM3_P0(TT2);
	}

	for (j = 16; j < 64; ++j) {
		T[j] = 0x7A879D8A;
		SS1 = SM3_ROTATELEFT((SM3_ROTATELEFT(A, 12) + E + SM3_ROTATELEFT(T[j], j)), 7);
		SS2 = SS1 ^ SM3_ROTATELEFT(A, 12);
		TT1 = SM3_FF1(A, B, C) + D + SS2 + W1[j];
		TT2 = SM3_GG1(E, F, G) + H + SS1 + W[j];
		D = C;
		C = SM3_ROTATELEFT(B, 9);
		B = A;
		A = TT1;
		H = G;
		G = SM3_ROTATELEFT(F, 19);
		F = E;
		E = SM3_P0(TT2);
	}

	digest[0] ^= A;
	digest[1] ^= B;
	digest[2] ^= C;
	digest[3] ^= D;
	digest[4] ^= E;
	digest[5] ^= F;
	digest[6] ^= G;
	digest[7] ^= H;
}

void sm3(const unsigned char* msg, size_t msglen,
	     unsigned char dgst[SM3_DIGEST_LENGTH])
{
	sm3_ctx_t ctx;

	sm3_init(&ctx);
	sm3_update(&ctx, msg, msglen);
	sm3_final(&ctx, dgst);

	memset(&ctx, 0, sizeof(sm3_ctx_t));
}